#!/usr/bin/env python3

# Copyright (c) 2025 Huawei Device Co., Ltd. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
import cv2
import numpy as np
import os
from datetime import datetime, timezone

class AppleDetector(Node):
    def __init__(self):
        super().__init__('apple_detector')
        self.get_logger().info("!!! 节点初始化完成 !!!")  
        self.get_logger().info("ROS日志系统初始化完成")  
        self.subscription = self.create_subscription(
            Image,
            'apple_image_topic',
            self.image_callback,
            10)
        
        # 设置结果保存路径
        self.result_dir = "apple_detection_results"
        os.makedirs(self.result_dir, exist_ok=True)
        self.get_logger().info(f"Apple detector ready, results will be saved to {self.result_dir}")

    def object_detect(image):
        hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)                               # 图像从BGR颜色模型转换为HSV模型
        mask_red = cv2.inRange(hsv_img, lower_red, upper_red)                           # pyright: ignore[reportUndefinedVariable]

        contours, hierarchy = cv2.findContours(mask_red, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE) 

        for cnt in contours:                                                          
            if cnt.shape[0] < 150:
                continue
                
            (x, y, w, h) = cv2.boundingRect(cnt)                                      
            cv2.drawContours(image, [cnt], -1, (0, 255, 0), 2)                        
            cv2.circle(image, (int(x+w/2), int(y+h/2)), 5, (0, 255, 0), -1)           
            
        cv2.imshow("object", image)                                                    
        cv2.waitKey(0)
        cv2.destroyAllWindows()     
    
    def image_callback(self, msg):
        try:
            # 将ROS图像消息转换为OpenCV格式
            if msg.encoding != 'bgr8':
                raise ValueError(f"Unsupported encoding: {msg.encoding}")
            
            image = np.frombuffer(msg.data, dtype=np.uint8)
            image = image.reshape((msg.height, msg.width, 3))
            
            # 苹果检测逻辑
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
            lower_red = np.array([0, 90, 128])     
            upper_red = np.array([180, 255, 255])  
            mask = cv2.inRange(hsv, lower_red, upper_red)
            
            # 查找轮廓并绘制边框
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            result_image = image.copy()  

            for cnt in contours:
                if cnt.shape[0] < 150:  
                    continue
                
                (x, y, w, h) = cv2.boundingRect(cnt)
                # 在 result_image 上绘制轮廓和中心点
                cv2.drawContours(result_image, [cnt], -1, (0, 255, 0), 2)  
                cv2.circle(result_image, (int(x+w/2), int(y+h/2)), 5, (0, 255, 0), -1)  
            
            # 保存结果
            timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d_%H%M%S")
            filename = os.path.join(self.result_dir, f"result_{timestamp}.jpg")
            cv2.imwrite(filename, result_image)  
            self.get_logger().info(f"Detection result saved: {filename}")
            
        except Exception as e:
            self.get_logger().error(f"Processing error: {str(e)}")

    

def main(args=None):
    rclpy.init(args=args)
    detector = AppleDetector()
    rclpy.spin(detector)
    detector.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()


